// 预先版本只考虑逐元素处理
#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2; 
constexpr int32_t CORE_NUM = 8; 

class KernelTraceCustom{
public:
    __aicore__ inline KernelTraceCustom() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,const TraceCustomTilingData &tilingData)
    {
        this->tiling  = tilingData;

        uint32_t blockIdx = AscendC::GetBlockIdx();
        uint32_t blockNum = AscendC::GetBlockNum();
        bool isLast = (tiling.fullTileLength != 0) && (blockIdx == blockNum - 1);
        matrixOrder = (tiling.fullTileLength == 0)
                      ? tiling.lastTileLength
                      : (isLast ? tiling.lastTileLength : tiling.fullTileLength);
        xGm.SetGlobalBuffer((__gm__ half *)x, tiling.matrixNum * tiling.rowLength * tiling.columnLength);
        yGm.SetGlobalBuffer((__gm__ half *)y, tiling.matrixNum);
        pipe.InitBuffer(inQueueDiag, BUFFER_NUM, tiling.fullTileLength * 32); 
        pipe.InitBuffer(outQueueY, BUFFER_NUM, sizeof(half));
        pipe.InitBuffer(workQueue, BUFFER_NUM, tiling.fullTileLength * 32);    
    }
    __aicore__ inline void Process()
    {
        CopyIn();     
        Compute();     
        CopyOut();    
    }

private:
    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<half> diagLocal = inQueueDiag.AllocTensor<half>();
        uint64_t matrixBaseOffset = (uint64_t)AscendC::GetBlockIdx() * tiling.fullTileLength * tiling.columnLength;
        uint64_t blockOffset = (uint64_t)AscendC::GetBlockIdx() * tiling.fullTileLength;
        uint64_t globalOffset = matrixBaseOffset + blockOffset;
        AscendC::DataCopyExtParams copyParams{
            static_cast<uint16_t>(matrixOrder),   
            static_cast<uint32_t>(sizeof(half)), 
            static_cast<uint32_t>((tiling.columnLength) * sizeof(half)),
            0,
            0};
        AscendC::DataCopyPadExtParams<half> padParams{true, 0, static_cast<uint8_t>(tiling.alignNum - 1), 0}; 
        AscendC::DataCopyPad(diagLocal, xGm[globalOffset], copyParams, padParams);
        inQueueDiag.EnQue(diagLocal);
    }
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<half> diagLocal    = inQueueDiag.DeQue<half>();
        AscendC::LocalTensor<half> yLocal       = outQueueY.AllocTensor<half>();
        AscendC::LocalTensor<half> workLocal   = workQueue.AllocTensor<half>();
        AscendC::ReduceSum<half>(yLocal, diagLocal, workLocal, matrixOrder * tiling.alignNum);
        outQueueY.EnQue<half>(yLocal);
        inQueueDiag.FreeTensor(diagLocal);
        workQueue.FreeTensor(workLocal);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<half> yLocal = outQueueY.DeQue<half>();
        AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>(0);
        AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(0);
        AscendC::SetAtomicAdd<half>();
        AscendC::DataCopy(yGm, yLocal, 32);
        AscendC::SetAtomicNone();
        outQueueY.FreeTensor(yLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueDiag;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> workQueue; 
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueY;
    AscendC::GlobalTensor<half> xGm;
    AscendC::GlobalTensor<half> yGm;

    TraceCustomTilingData tiling;
    uint32_t matrixNum;            // 矩阵数量
    uint32_t rowLength;            // 矩阵行数
    uint32_t columnLength;         // 矩阵列数
    uint32_t fullTileLength;       // 主核处理的元素个数
    uint32_t lastTileLength;       // 尾核处理的元素个数
    uint32_t alignPadding;         // 对齐填充长度
    uint32_t matrixOrder;          // 当前核处理的矩阵阶数（fullTileLength或lastTileLength）

};

extern "C" __global__ __aicore__ void trace_custom(GM_ADDR x,  GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    (void)workspace;
    GET_TILING_DATA(tilingData, tiling);
    KernelTraceCustom op;
    op.Init(x, y, tilingData);
    op.Process();
}